package com.liguangting.agentsflex.simple;

import com.agentsflex.core.document.Document;
import com.agentsflex.core.llm.Llm;
import com.agentsflex.core.store.DocumentStore;
import com.agentsflex.core.store.SearchWrapper;
import com.agentsflex.llm.qwen.QwenLlm;
import com.agentsflex.llm.qwen.QwenLlmConfig;
import com.agentsflex.store.aliyun.AliyunVectorStore;
import com.agentsflex.store.aliyun.AliyunVectorStoreConfig;
import com.liguangting.agentsflex.constant.ApiKeys;

import java.util.List;


/**
 * @author liguangting
 * @date 2024/7/5
 * @description
 */
public class AliVectorStoreTest {

    public static void main(String[] args) {
        AliyunVectorStoreConfig storeConfig = new AliyunVectorStoreConfig();

        //设置阿里云向量数据检索服务的相关配置
        storeConfig.setApiKey(ApiKeys.ALI_VECTOR_DB_API_KEY);
        storeConfig.setEndpoint(ApiKeys.ALI_VECTOR_DB_ENDPOINT);
        storeConfig.setDefaultCollectionName(ApiKeys.ALI_VECTOR_DB_COLLECTION);

        DocumentStore store = new AliyunVectorStore(storeConfig);

        //创建 Embedding 模型
        QwenLlmConfig config = new QwenLlmConfig();
        config.setApiKey(ApiKeys.QWEN_API_KEY);
        Llm llm = new QwenLlm(config);

        //为 store 配置 Embedding 模型
        store.setEmbeddingModel(llm);

        // 新增数据
        Document document = new Document();
        document.setId(10000);
        document.setContent("2024年7月5日，杭州西湖边的一家咖啡店，一位女士正在欣赏着窗外的风景。她说，这里的环境很好，很安静，适合工作和学习。");
        store.store(document);

        // 数据召回
        SearchWrapper wrapper = new SearchWrapper();
        wrapper.setText("杭州西湖边咖啡店");
        List<Document> result = store.search(wrapper);
        System.out.println(result);
    }
}
